"""
遍历某路径下的所有子文件夹，分别计算所有子文件夹中TIFF图像的总值栅格
需要所有的TIFF图像空间参考、地理转换参数一致
"""
from osgeo import gdal
import numpy as np
import os

def calculate_monthly_average(input_parent_folder, output_parent_folder):
    # 遍历1到12月的子文件夹
    for month in range(2002, 2021):
        # 格式化子文件夹名称，例如 "01", "02", ..., "12"
        month_folder = f"{month:04d}"
        input_folder = os.path.join(input_parent_folder, month_folder)
        #注意根据待处理数据集的名称更改此输出文件名
        output_file = os.path.join(output_parent_folder, f"CO2_emission_{month_folder}.tif")

        # 如果月份文件夹不存在，跳过
        if not os.path.exists(input_folder):
            print(f"文件夹 {input_folder} 不存在，跳过...")
            continue

        print(f"正在处理 {month_folder} 月的数据...")

        arr_mean = None
        count = 0
        geotransform = None
        projection = None

        # 遍历该月份文件夹中的所有 TIFF 文件
        for file in os.listdir(input_folder):
            if file.endswith(".tif"):
                tiff_path = os.path.join(input_folder, file)
                dataset = gdal.Open(tiff_path)
                band = dataset.GetRasterBand(1)
                arr = band.ReadAsArray()

                # 初始化累加数组
                if arr_mean is None:
                    arr_mean = np.zeros_like(arr, dtype=float)

                # 累加每个栅格
                arr_mean += arr
                count += 1

                # 记录地理信息
                if geotransform is None:
                    geotransform = dataset.GetGeoTransform()
                if projection is None:
                    projection = dataset.GetProjection()

        # 如果没有文件，跳过该月份
        if count == 0:
            print(f"{month_folder} 月没有有效的 TIFF 文件，跳过...")
            continue

        # 计算平均值
        # arr_mean /= count

        # 保存平均(或总和)栅格为 TIFF 文件
        nrows, ncols = arr_mean.shape
        driver = gdal.GetDriverByName("GTiff")
        output_dataset = driver.Create(output_file, ncols, nrows, 1, gdal.GDT_Float32)
        output_dataset.SetGeoTransform(geotransform)
        output_dataset.SetProjection(projection)

        output_band = output_dataset.GetRasterBand(1)
        output_band.WriteArray(arr_mean)
        output_band.FlushCache()
        output_band = None

        print(f"{month_folder} 月平均 TIFF 文件已保存: {output_file}")

# 主函数
def main():
    input_parent_folder = r"D:\Lenovo\Desktop\tmp\CO2_tiff"
    output_parent_folder = r"D:\Lenovo\Desktop\tmp\CO2_tiff\yearly_CO2_emission"

    # 如果输出父文件夹不存在，则创建
    if not os.path.exists(output_parent_folder):
        os.makedirs(output_parent_folder)

    # 调用计算函数
    calculate_monthly_average(input_parent_folder, output_parent_folder)

main()
